import os
import random
import shutil
from collections import defaultdict
from pathlib import Path

def get_class_prefix(filename: str) -> str:
    """
    从文件名中提取类别前缀。
    例如: 'Dessert01_40.jpg' -> 'Dessert01'
    如果您的命名规则不同（例如使用'-'而不是'_'），请修改此函数。
    """
    # 以第一个下划线为准进行分割
    return filename.split('_')[0]

def split_dataset_by_class(source_images_dir: str, source_labels_dir: str, dest_base_dir: str, val_split_ratio: float = 0.2):
    """
    根据文件名中的类别前缀，将数据集按比例划分为训练集和验证集。

    Args:
        source_images_dir (str): 包含所有原始图片的文件夹路径。
        source_labels_dir (str): 包含所有原始标签的文件夹路径。
        dest_base_dir (str): 目标根目录。脚本将在此目录下创建 'images' 和 'labels' 文件夹。
        val_split_ratio (float): 分配给验证集的数据比例。
    """
    # --- 1. 准备路径和创建目标文件夹 ---
    source_images_path = Path(source_images_dir)
    source_labels_path = Path(source_labels_dir)
    dest_base_path = Path(dest_base_dir)

    print("正在创建目标文件夹...")
    for folder in ['images', 'labels']:
        for split in ['train', 'val']:
            (dest_base_path / folder / split).mkdir(parents=True, exist_ok=True)
    print(f"目标文件夹结构已在 '{dest_base_path}' 中创建完毕。")
    print("-" * 40)

    # --- 2. 按类别对文件进行分组 ---
    # defaultdict(list) 让我们能方便地向字典的值（列表）中添加元素
    files_by_class = defaultdict(list)
    
    # 我们只需要遍历图片文件夹，因为标签文件是与之对应的
    try:
        image_files = [f for f in os.listdir(source_images_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    except FileNotFoundError:
        print(f"错误：找不到源图片文件夹 '{source_images_path}'。请检查路径是否正确。")
        return

    if not image_files:
        print(f"错误：在 '{source_images_path}' 中未找到任何图片文件。")
        return

    # 将所有图片按类别前缀分组
    for img_filename in image_files:
        class_prefix = get_class_prefix(img_filename)
        files_by_class[class_prefix].append(img_filename)
    
    print(f"扫描完成！共找到 {len(files_by_class)} 个类别。")
    print("每个类别的统计数量如下：")
    for class_name, files in files_by_class.items():
        print(f"  - 类别 '{class_name}': {len(files)} 个文件")
    print("-" * 40)

    # --- 3. 为每个类别划分训练集和验证集并复制文件 ---
    total_train_count = 0
    total_val_count = 0

    print("开始划分数据集并复制文件...")
    for class_name, filenames in files_by_class.items():
        # 随机打乱当前类别的所有文件
        random.shuffle(filenames)
        
        # 计算验证集的数量
        num_val = int(len(filenames) * val_split_ratio)
        
        # 进行分割
        val_files = filenames[:num_val]
        train_files = filenames[num_val:]

        print(f"处理类别 '{class_name}': {len(train_files)} 个用于训练, {len(val_files)} 个用于验证。")
        total_train_count += len(train_files)
        total_val_count += len(val_files)

        # 复制验证集文件
        for filename in val_files:
            basename, _ = os.path.splitext(filename)
            # 复制图片
            shutil.copy(
                source_images_path / filename,
                dest_base_path / 'images' / 'val' / filename
            )
            # 复制对应的标签
            shutil.copy(
                source_labels_path / f"{basename}.txt",
                dest_base_path / 'labels' / 'val' / f"{basename}.txt"
            )

        # 复制训练集文件
        for filename in train_files:
            basename, _ = os.path.splitext(filename)
            # 复制图片
            shutil.copy(
                source_images_path / filename,
                dest_base_path / 'images' / 'train' / filename
            )
            # 复制对应的标签
            shutil.copy(
                source_labels_path / f"{basename}.txt",
                dest_base_path / 'labels' / 'train' / f"{basename}.txt"
            )

    print("-" * 40)
    print("处理完成！")
    print(f"总计: {total_train_count} 张训练图片, {total_val_count} 张验证图片。")
    print(f"所有文件已保存至 '{dest_base_path}' 目录。")


if __name__ == '__main__':
    # ===================================================================
    #                           请在这里配置您的路径
    # ===================================================================
    
    # 1. 存放您所有原始图片的文件夹路径
    # 例如: 'C:/Users/YourUser/Desktop/My_Raw_Data/images'
    SOURCE_IMAGES_DIR = '/home/youngqi/Documents/datasets/datasets_all/images' 
    
    # 2. 存放您所有原始标签 (.txt) 的文件夹路径
    # 例如: 'C:/Users/YourUser/Desktop/My_Raw_Data/labels'
    SOURCE_LABELS_DIR = '/home/youngqi/Documents/datasets/datasets_all/labels'

    # 3. 您希望将新生成的数据集存放到的目标文件夹
    # 脚本会自动在此文件夹下创建 'images/train', 'images/val' 等结构
    # 例如: 'C:/Users/YourUser/Desktop/coco_dataset'
    DESTINATION_BASE_DIR = '/home/youngqi/Downloads/datasheet/'

    # 4. 验证集所占的比例（例如 0.2 代表 20%）
    VALIDATION_SPLIT_RATIO = 0.2

    # ===================================================================
    #                           配置结束
    # ===================================================================
    
    # 运行主函数
    split_dataset_by_class(
        SOURCE_IMAGES_DIR,
        SOURCE_LABELS_DIR,
        DESTINATION_BASE_DIR,
        VALIDATION_SPLIT_RATIO
    )